package main

import (
	"strings"
	"sync"
)

const IllegalAccessMaxCount int = 3

type IpBlacklist struct {
	datas map[string]int
	mutex sync.Mutex
}

func NewIpBlacklist() *IpBlacklist {
	ibl := IpBlacklist{datas: make(map[string]int)}
	return &ibl
}

func (p *IpBlacklist) Add(ip string) int {
	p.mutex.Lock()
	defer p.mutex.Unlock()

	ip = p.TrimIp(ip)

	if v, ok := p.datas[ip]; ok {
		v++
		p.datas[ip] = v
		return v
	} else {
		p.datas[ip] = 1
		return 1
	}
}

func (p *IpBlacklist) Find(ip string) int {
	ip = p.TrimIp(ip)

	if v, ok := p.datas[ip]; ok {
		return v
	} else {
		return 0
	}
}

type ImportData struct {
	Ip    string
	Count int
}

func (p *IpBlacklist) Import(datas []ImportData) {
	p.mutex.Lock()
	defer p.mutex.Unlock()

	p.datas = make(map[string]int)
	for _, v := range datas {
		p.datas[v.Ip] = v.Count
	}
}

func (p *IpBlacklist) IsLegal(ip string) bool {
	ip = p.TrimIp(ip)

	if v, ok := p.datas[ip]; ok {
		if v >= IllegalAccessMaxCount {
			return false
		}
	}

	return true
}

func (p *IpBlacklist) TrimIp(ip string) string {
	if index := strings.Index(ip, ":"); index > 0 {
		ip = ip[:index]
	}
	return ip
}
